package cn.lsh.kafka.partitioner;

import org.apache.kafka.clients.consumer.internals.AbstractPartitionAssignor;
import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.utils.CircularIterator;
import org.apache.kafka.common.utils.Utils;

import java.util.*;
import java.util.stream.Collectors;

/**
 * 自定义消费端分区器
 */
public class MyConsumerAssignor extends AbstractPartitionAssignor {

	@Override
	public Map<String, List<TopicPartition>> assign(Map<String, Integer> partitionsPerTopic, Map<String, Subscription> subscriptions) {
		Map<String, List<TopicPartition>> assignment = new HashMap();
		List<MemberInfo> memberInfoList = new ArrayList();
		for (Map.Entry<String, Subscription> entry : subscriptions.entrySet()) {
			assignment.put(entry.getKey(), new ArrayList<>());
			memberInfoList.add(new MemberInfo(entry.getKey(), entry.getValue().groupInstanceId()));
		}
		CircularIterator<MemberInfo> assigner = new CircularIterator(Utils.sorted(memberInfoList));
		List<TopicPartition> topicPartitions = allPartitionsSorted(partitionsPerTopic, subscriptions);
		topicPartitions.forEach(topicPartition -> {
			if (subscriptions.get(assigner.peek().memberId).topics().contains(topicPartition.topic())) {
				assigner.next();
			}
			assignment.get(assigner.next().memberId).add(topicPartition);
		});
		return assignment;
	}

	private List<TopicPartition> allPartitionsSorted(Map<String, Integer> partitionsPerTopic, Map<String, Subscription> subscriptions) {
		//订阅信息
		List<TopicPartition> allPartitions = new ArrayList();
		//获取所有的订阅的topics
		Set<String> topics = subscriptions.values().stream().flatMap(s -> s.topics().stream()).collect(Collectors.toCollection(TreeSet::new));
		topics.forEach(topic -> {
			Integer numPartitionsForTopic = partitionsPerTopic.get(topic);
			if (numPartitionsForTopic != null) {
				allPartitions.addAll(AbstractPartitionAssignor.partitions(topic, numPartitionsForTopic));
			}
		});
		return allPartitions;
	}

	@Override
	public String name() {
		return "my_assignor";
	}

}
